#!/usr/bin/env python
import itk
import sys
from optparse import OptionParser
import os
import imageTools.ITKUtils.io as io
import imageTools.ITKUtils.ifilter as ifilter
import numpy as np

def getParser():
    parser = OptionParser()
    parser.add_option("-i","--img",dest='fname',
                      help='name of image to read')
    parser.add_option("-t","--thresh",dest='threshold',nargs=2,default=(0,10000),type='int')
    parser.add_option("-s","--seed",dest='seed',nargs=3,default=(0,0,0),
                      type='int')
    parser.add_option("-m", "--mask", dest="maskfile" , default = 'None' )
    return parser

def run_seg(fname, maskfile, threshold, seed):
    rimg = io.readImage(fname, imgMode = 'float')
    rimg_arr = itk.PyBuffer.IF3.GetArrayFromImage(rimg)

    if maskfile != 'None' : 
         mask = io.readImage(maskfile, imgMode = 'float')
         mask_arr = itk.PyBuffer.IF3.GetArrayFromImage(mask)
         rimg_arr = rimg_arr*mask_arr
    rimg = itk.PyBuffer.IF3.GetImageFromArray(rimg_arr) 
    filter = itk.ConnectedThresholdImageFilter.IF3IF3.New()
    filter.SetInput(rimg)
    filter.SetLower(threshold[0])
    filter.SetUpper(threshold[1])
    filter.AddSeed(seed)
    mask = filter.GetOutput()
    mask.Update()
    # get output name
    tmp = os.path.splitext(fname)
    io.writeImage(mask,tmp[0]+"_seg"+tmp[1],outMode="ushort")
def main():
    parser = getParser()
    (options, args) = parser.parse_args()
    print options
    run_seg(options.fname, options.maskfile, options.threshold, options.seed)
      
    
if __name__ == '__main__':
    main()


